{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the Spherical harmonics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import paddle\n",
    "\n",
    "from paddle_harmonics.quadrature import legendre_gauss_weights\n",
    "from paddle_harmonics.legendre import legpoly, clm\n",
    "from paddle_harmonics import RealSHT, InverseRealSHT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forming the Vandermonde matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlat = 60\n",
    "nlon = 2*nlat\n",
    "lmax = mmax = nlat\n",
    "\n",
    "sht = RealSHT(nlat, nlon, lmax=lmax, mmax=mmax)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "\n",
    "nmodes = int(lmax * (lmax+1) / 2)\n",
    "e = paddle.zeros([nmodes, lmax, mmax], dtype=paddle.complex64)\n",
    "\n",
    "midx = lambda l,m : l*(l+1)//2 + m\n",
    "\n",
    "for l in range(lmax):\n",
    "    for m in range(l+1):\n",
    "        e[midx(l,m), l, m] = 1.0\n",
    "\n",
    "vdm = isht(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting the Spherical Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_lmax = 6\n",
    "\n",
    "fig = plt.figure(layout='constrained', figsize=(12, 8))\n",
    "subfigs = fig.subfigures(plt_lmax, plt_lmax)\n",
    "\n",
    "for l in range(plt_lmax):\n",
    "    for m in range(l+1):\n",
    "        plot_sphere(vdm[midx(l,m)], fig=subfigs[l, m], title=f\"l={l}, m={m}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's also make an animated version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.animation as animation\n",
    "\n",
    "fig = plt.figure(layout='constrained', figsize=(12, 8), dpi=72)\n",
    "subfigs = fig.subfigures(plt_lmax, plt_lmax)\n",
    "\n",
    "moviewriter = animation.writers['pillow'](fps=30)\n",
    "moviewriter.setup(fig, 'spherical_harmonics.gif', dpi=72)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for frame in range(120):\n",
    "    # compute the rotation of the sphere\n",
    "    lon = -120 - 3 * frame\n",
    "\n",
    "    if frame % 10 == 0:\n",
    "        print(f\"frame={frame}\")\n",
    "    \n",
    "    for l in range(plt_lmax):\n",
    "        for m in range(l+1):\n",
    "            plot_sphere(vdm[midx(l,m)], fig=subfigs[l, m], title=f\"l={l}, m={m}\", central_longitude=lon)\n",
    "\n",
    "    plt.draw()\n",
    "    moviewriter.grab_frame()\n",
    "\n",
    "moviewriter.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
